package com.zhiziai.aiproxy.utils;

import cn.hutool.core.util.RandomUtil;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.google.common.base.Preconditions;
import com.zhiziai.aiproxy.dto.CommonDto;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import okhttp3.MediaType;
import okhttp3.MultipartBody;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;
import org.apache.commons.lang3.StringUtils;
import org.apache.tika.Tika;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.TimeUnit;

import static cn.hutool.http.HttpUtil.getMimeType;

@Slf4j
public class OpenAIUtil {

    static String CHAT_URL = "https://api.openai.com/v1/chat/completions";
    static String DALLE_GEN_URL = "https://api.openai.com/v1/images/generations";
    static String DALLE_EDIT_URL = "https://api.openai.com/v1/images/edits";
    static String DALLE_VARIATIONS_URL = "https://api.openai.com/v1/images/variations";
    static String AUDIO_TRANSCRIPTIONS_URL = "https://api.openai.com/v1/audio/transcriptions";
    static List<String> APIKEYS = Arrays.asList("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");
    static String CHATGPT_MODEL_VERSION = "gpt-3.5-turbo";
    static String WHISPER_MODEL_VERSION = "whisper-1";

    @Nullable
    public static Msg chatgpt(String question) throws Exception{
        log.info("question {} " ,question);
        long start = System.currentTimeMillis();
        Msg resMsg = null;
        Preconditions.checkArgument(StringUtils.isNotEmpty(question), "question is null");
        Msg reqMsg = Msg.builder().content(question).role("user").build();
        JSONObject req = new JSONObject();
        req.put("model", CHATGPT_MODEL_VERSION);
        req.put("messages", Arrays.asList(reqMsg));
        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body = RequestBody.create(mediaType, JSONObject.toJSONString(req));
        CommonDto<Integer,String> responseBody = post(CHAT_URL,body,1800);
        Preconditions.checkArgument(responseBody.getId() == 200, responseBody.getRes());
        log.info("success answer is {} , cost {} " ,responseBody, System.currentTimeMillis() - start);
        JSONObject res = JSONObject.parseObject(responseBody.getRes());
        JSONArray jsonArray = res.getJSONArray("choices");
        if(jsonArray.size() > 0){
            JSONObject choice = jsonArray.getJSONObject(0);
            resMsg = choice.getObject("message",Msg.class);
        }
        return resMsg;
    }

    public static List<String> dalleGen(String prompt, Integer genNum, String resulotionRate) throws Exception{
        log.info("prompt {} " ,prompt);
        long start = System.currentTimeMillis();
        List<String> urls = new ArrayList<>();
        JSONObject req = new JSONObject();
        req.put("prompt", prompt);
        req.put("n", genNum);
        req.put("size", resulotionRate);
        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body = RequestBody.create(mediaType, JSONObject.toJSONString(req));
        CommonDto<Integer,String> responseBody = post(DALLE_GEN_URL,body,1800);
        Preconditions.checkArgument(responseBody.getId() == 200, responseBody.getRes());
        log.info("success answer is {}, cost {} " ,responseBody, System.currentTimeMillis() - start);
        JSONObject res = JSONObject.parseObject(responseBody.getRes());
        JSONArray jsonArray = res.getJSONArray("data");
        if(jsonArray.size() > 0){
            for (int i = 0; i < jsonArray.size(); i++) {
                JSONObject urlJo = jsonArray.getJSONObject(i);
                String url = urlJo.getString("url");
                urls.add(url);
            }
            return urls;
        }
        return null;
    }

    public static List<String> dalleEdit(String prompt, int genNum, File image, File mask, String resulotionRate) throws Exception{
        List<String> urls = new ArrayList<>();
        String imgFileType = getMimeType(image.getName());
        RequestBody imageBody = RequestBody.create(MediaType.parse(imgFileType), image);
        String imageName = image.getName();

//        if(mask != null){
//            String maskFileType = getMimeType(mask.getName());
//            RequestBody maskBody = RequestBody.create(MediaType.parse(maskFileType), mask);
//            String maskName = mask.getName();
//        }

        //这边是把file写进来，也有写路径的，但我这边是写file文件，parse不行的话可以直接改这个"multipart/form-data"
        MultipartBody reqBody = new MultipartBody.Builder()
                .setType(MultipartBody.FORM) // 设置传参为form-data格式
                .addFormDataPart("prompt", prompt)
                .addFormDataPart("n", genNum+"")
                .addFormDataPart("size", resulotionRate)
                .addFormDataPart("image", imageName, imageBody) // 中间参数为文件名
//                .addFormDataPart("mask", maskName, maskBody) // 中间参数为文件名
                .build();
        CommonDto<Integer,String> responseBody = post(DALLE_EDIT_URL,reqBody,1800);
        Preconditions.checkArgument(responseBody.getId() == 200, responseBody.getRes());
        log.info("success answer is {}" ,responseBody);
        JSONObject res = JSONObject.parseObject(responseBody.getRes());
        JSONArray jsonArray = res.getJSONArray("data");
        if(jsonArray.size() > 0){
            for (int i = 0; i < jsonArray.size(); i++) {
                JSONObject urlJo = jsonArray.getJSONObject(i);
                String url = urlJo.getString("url");
                urls.add(url);
            }
            return urls;
        }
        return null;
    }

    public static List<String> dalleVariations(int genNum, File image, String resulotionRate) throws Exception{
        List<String> urls = new ArrayList<>();
        String imgFileType = getMimeType(image.getName());
        RequestBody imageBody = RequestBody.create(MediaType.parse(imgFileType), image);
        String imageName = image.getName();

        //这边是把file写进来，也有写路径的，但我这边是写file文件，parse不行的话可以直接改这个"multipart/form-data"
        MultipartBody reqBody = new MultipartBody.Builder()
                .setType(MultipartBody.FORM) // 设置传参为form-data格式
                .addFormDataPart("n", genNum+"")
                .addFormDataPart("size", resulotionRate)
                .addFormDataPart("image", imageName, imageBody) // 中间参数为文件名
                .build();
        CommonDto<Integer,String> responseBody = post(DALLE_VARIATIONS_URL,reqBody,1800);
        Preconditions.checkArgument(responseBody.getId() == 200, responseBody.getRes());
        log.info("success answer is {}" ,responseBody);
        JSONObject res = JSONObject.parseObject(responseBody.getRes());
        JSONArray jsonArray = res.getJSONArray("data");
        if(jsonArray.size() > 0){
            for (int i = 0; i < jsonArray.size(); i++) {
                JSONObject urlJo = jsonArray.getJSONObject(i);
                String url = urlJo.getString("url");
                urls.add(url);
            }
            return urls;
        }
        return null;
    }

    @Nullable
    public static String audioTranscriptionsEncrypt(MultipartFile multipartFile,String language,String prompt,String temperature) throws Exception{
        File file = FileUtil.transferToFile(multipartFile);
        Tika tika = new Tika();
        String fileMimeType = tika.detect(file);
        RequestBody fileBody = RequestBody.create(MediaType.parse(fileMimeType), file);
        String fileName = file.getName();
        MultipartBody reqBody = new MultipartBody.Builder()
                .setType(MultipartBody.FORM) // 设置传参为form-data格式
                .addFormDataPart("file", fileName, fileBody)
                .addFormDataPart("model",WHISPER_MODEL_VERSION)
                .addFormDataPart("language",language)
                .addFormDataPart("prompt",prompt)
                .addFormDataPart("temperature",temperature)
                .build();
        CommonDto<Integer,String> responseBody = post(AUDIO_TRANSCRIPTIONS_URL,reqBody,1800);
        Preconditions.checkArgument(responseBody.getId() == 200, responseBody.getRes());
        JSONObject answerJo = JSONObject.parseObject(responseBody.getRes());
        String answerText = answerJo.getString("text");
        log.info("success answer is {}" ,answerText);
        return answerText;
    }

    public static CommonDto<Integer,String> post(String url , RequestBody body, Integer timeOut) throws Exception{
        int retryCount = 0;
        CommonDto<Integer,String> res = new CommonDto<>();
        while (retryCount < 2) {
            Response response = okCall(url, body, timeOut);
            int resCode = response.code();
            String resBody = response.body() != null ? response.body().string() : null;
            res.setRes(resBody);
            res.setId(resCode);
            if(resCode == 200){
                break;
            }else{
                retryCount++;
                log.info("post error {} , retryCount {}", res,retryCount);
            }
        }
        return res;
    }

    @NotNull
    private static Response okCall(String url, RequestBody body, Integer timeOut) throws IOException {
        OkHttpClient client = new OkHttpClient().newBuilder()
                .connectTimeout(timeOut, TimeUnit.SECONDS)
                .writeTimeout(timeOut,TimeUnit.SECONDS)
                .readTimeout(timeOut,TimeUnit.SECONDS)
                .build();
        Request request = new Request.Builder()
                .url(url)
                .post(body)
                .addHeader("Authorization", "Bearer " + RandomUtil.randomEle(APIKEYS))
//                .addHeader("Content-Type", "application/json")
                .build();
        return client.newCall(request).execute();
    }


    public static String testChatGpt(String question) throws Exception{
        Preconditions.checkArgument(StringUtils.isNotEmpty(question), "question is null");
        Msg reqMsg = Msg.builder().content(question).role("user").build();
        JSONObject req = new JSONObject();
        req.put("model", CHATGPT_MODEL_VERSION);
        req.put("messages", Arrays.asList(reqMsg));
        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body = RequestBody.create(mediaType, JSONObject.toJSONString(req));
        OkHttpClient client = new OkHttpClient().newBuilder()
                .connectTimeout(1800, TimeUnit.SECONDS)
                .writeTimeout(1800,TimeUnit.SECONDS)
                .readTimeout(1800,TimeUnit.SECONDS)
                .build();
        Request request = new Request.Builder()
                .url(CHAT_URL)
                .post(body)
                .addHeader("Authorization", "Bearer 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")
//                .addHeader("Content-Type", "application/json")
                .build();
        Response response = client.newCall(request).execute();
        int resCode = response.code();
        String resBody = response.body() != null ? response.body().string() : null;
        Preconditions.checkArgument(resCode == 200, resBody);
        log.info("success test question {} , answer is {}" ,reqMsg,resBody);
        JSONObject res = JSONObject.parseObject(resBody);
        JSONArray jsonArray = res.getJSONArray("choices");
        if(jsonArray.size() > 0){
            JSONObject choice = jsonArray.getJSONObject(0);
            return choice.getObject("message",Msg.class).getContent();
        }
        return "";
    }

    public static void main(String[] args) throws Exception{
        String apiKey = RandomUtil.randomEle(APIKEYS);
        System.out.println(apiKey);
        chatgpt("123123");

    }


    @Data
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    public static class Msg{
        private String role;
        private String content;
    }
}

